package com.dzsw.datastruct.tree.bst;

import com.dzsw.datastruct.tree.binary.IterationBinaryTreeImpl;
import com.dzsw.datastruct.tree.binary.Node;

import java.util.Comparator;

public class BinarySearchTree<E> extends IterationBinaryTreeImpl<E> {

    private final Comparator<E> comparator;

    public BinarySearchTree() {
        this(null);
    }

    public BinarySearchTree(final Comparator<E> comparator) {
        this.comparator = comparator;
    }

    /**
     * 增加节点
     *
     * @param element
     */
    public void add(E element) {
        elementNotNullCheck(element);
        // 空树，添加第一个节点
        if (root == null) {
            root = Node.createNode(element, null);
            size++;
            return;
        }
        // 非空树情况，找父节点
        Node<E> node = this.root;
        // 保存找到的父节点，默认设置为根节点
        Node<E> parent = node;
        // 找到父节点
        int result = 0;
        while (node != null) {
            result = compare(element, node.getElement());
            if (result > 0) {
                // 大于父节点值，取右子节点比较
                parent = node;
                node = node.getRight();
            } else if (result < 0) {
                parent = node;
                // 小于父节点值，取左子节点比较
                node = node.getLeft();
            } else {
                // 相等，返回不处理
                return;
            }
        }
        // 插入新节点
        Node<E> newNode = Node.createNode(element, parent);
        if (result > 0) {
            parent.setRight(newNode);
        } else {
            parent.setLeft(newNode);
        }
    }

    /**
     * 判断树是否包含值为 element 的节点
     *
     * @param element
     * @return
     */
    public boolean contains(E element) {
        return searchByElement(element) != null;
    }

    /**
     * 删除元素为element的节点
     *
     * @param element
     */
    public void remove(E element) {
        remove(searchByElement(element));
    }

    /**
     * 删除传入的节点
     *
     * @param node
     */
    private void remove(Node<E> node) {
        if (node == null) {
            return;
        }
        size--;
        // 删除度为 2 的节点，实际上转化为删除度为 1 或者 0 的节点。
        if (node.isFullNodes()) {
            // 找到后继节点
            Node<E> snode = this.successor(node);
            // 用后继几点的值覆盖度为 2 的节点的值
            node.setElement(snode.getElement());
            // 删除后继节点
            node = snode;
        }
        // 删除 node 节点（ node 的度必然是 1 或者 0 )
        Node<E> replacement = node.getLeft() != null ? node.getLeft() : node.getRight();
        // node 是度为 1 的节点
        if (replacement != null) {
            // 更改 parent
            replacement.setParent(node.getParent());
            // 更改 parent 的 left 、 right 的指向
            if (node.getParent() == null) {
                // node 是度为 1 的节点并且是根节点
                this.root = replacement;
            } else if (node == node.getParent().getLeft()) {
                node.getParent().setLeft(replacement);
            } else {
                // node == node.getParent().getRight()
                node.getParent().setRight(replacement);
            }
        } else if (node.getParent() == null) {
            // node 是叶子节点并且是根节点
            this.root = null;
        } else {
            // node 是叶子节点，但不是根节点
            if (node == node.getParent().getLeft()) {
                node.getParent().setLeft(null);
            } else {
                // node == node.getParent().getRight()
                node.getParent().setRight(null);
            }
        }
    }

    /**
     * 比较函数，返回0，e1==e2;返回值大于0，e1>e2;返回小于0，e1<e2
     *
     * @param e1
     * @param e2
     *
     * @return
     */
    private int compare(E e1, E e2) {
        if (comparator != null) {
            return comparator.compare(e1, e2);
        }
        return ((Comparable) e1).compareTo(e2);
    }

    /**
     * 新节点元素非空检查
     * @param element
     */
    private void elementNotNullCheck(E element) {
        if (element == null) {
            throw new IllegalArgumentException("element must not be null");
        }
    }

    /**
     * 查询值为 element 的节点
     *
     * @param element
     * @return
     */
    private Node<E> searchByElement(E element) {
        Node<E> node = this.root;
        int result = 0;
        while (node != null) {
            result = compare(element, node.getElement());
            if (result == 0) {
                return node;
            } else if (result > 0) {
                node = node.getRight();
            } else {
                node = node.getLeft();
            }
        }
        return null;
    }
}
